package websocket

import (
	"context"
	"errors"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var (
	// DefaultUpgrader 默认的 WebSocket Upgrader
	DefaultUpgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true // 允许所有来源，生产环境需要配置
		},
	}
)

// Connection WebSocket 连接封装
type Connection struct {
	conn        *websocket.Conn
	send        chan []byte
	mu          sync.RWMutex
	isClosed    bool
	messageType int
	pingPeriod  time.Duration
	pongWait    time.Duration
	writeWait   time.Duration
	ctx         context.Context
	cancel      context.CancelFunc
}

// Message WebSocket 消息
type Message struct {
	Type    string      `json:"type"`
	Payload interface{} `json:"payload"`
	Data    interface{} `json:"data,omitempty"`
}

// Hub 管理所有 WebSocket 连接
type Hub struct {
	// 注册的连接
	connections map[*Connection]bool

	// 从连接中接收消息
	broadcast chan []byte

	// 注册新连接
	register chan *Connection

	// 注销连接
	unregister chan *Connection

	// 连接ID映射
	connectionsByID map[string]*Connection

	mu sync.RWMutex
}

// NewHub 创建新的 Hub
func NewHub() *Hub {
	return &Hub{
		connections:     make(map[*Connection]bool),
		broadcast:       make(chan []byte, 256),
		register:        make(chan *Connection),
		unregister:      make(chan *Connection),
		connectionsByID: make(map[string]*Connection),
	}
}

// Run 运行 Hub（应在独立的 goroutine 中运行）
func (h *Hub) Run() {
	for {
		select {
		case conn := <-h.register:
			h.mu.Lock()
			h.connections[conn] = true
			h.mu.Unlock()

		case conn := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.connections[conn]; ok {
				delete(h.connections, conn)
				close(conn.send)
			}
			h.mu.Unlock()

		case message := <-h.broadcast:
			h.mu.RLock()
			for conn := range h.connections {
				select {
				case conn.send <- message:
				default:
					close(conn.send)
					delete(h.connections, conn)
				}
			}
			h.mu.RUnlock()
		}
	}
}

// Broadcast 广播消息给所有连接
func (h *Hub) Broadcast(message []byte) {
	h.broadcast <- message
}

// SendToID 发送消息给指定 ID 的连接
func (h *Hub) SendToID(id string, message []byte) error {
	h.mu.RLock()
	defer h.mu.RUnlock()

	conn, ok := h.connectionsByID[id]
	if !ok {
		return errors.New("connection not found")
	}

	select {
	case conn.send <- message:
		return nil
	default:
		return errors.New("connection send buffer full")
	}
}

// RegisterConnection 注册连接并关联 ID
func (h *Hub) RegisterConnection(id string, conn *Connection) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.connectionsByID[id] = conn
	h.connections[conn] = true
}

// UnregisterConnection 注销连接
func (h *Hub) UnregisterConnection(id string, conn *Connection) {
	h.mu.Lock()
	defer h.mu.Unlock()
	delete(h.connectionsByID, id)
	if _, ok := h.connections[conn]; ok {
		delete(h.connections, conn)
	}
}

// GetConnectionCount 获取连接数
func (h *Hub) GetConnectionCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.connections)
}

// GetConnections 获取所有连接ID
func (h *Hub) GetConnections() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	ids := make([]string, 0, len(h.connectionsByID))
	for id := range h.connectionsByID {
		ids = append(ids, id)
	}
	return ids
}
